package com.feizhaiyou.encrypt.codec;

import cn.hutool.core.codec.Base64Encoder;
import cn.hutool.core.util.CharsetUtil;
import cn.hutool.core.util.HexUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.crypto.SecureUtil;
import cn.hutool.crypto.SmUtil;
import cn.hutool.crypto.asymmetric.KeyType;
import cn.hutool.crypto.asymmetric.SM2;
import com.feizhaiyou.encrypt.constants.SecurityMode;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.security.KeyPair;
import java.util.HashMap;
import java.util.Map;

/**
 * @author ls
 */
@Data
@Slf4j
public class SM2Processor implements SecurityProcessor {

    private String publicKey;

    private String privateKey;

    private SM2 sm2;

    public SM2Processor(String publicKey, String privateKey) {
        if (StringUtils.isBlank(publicKey) || StringUtils.isBlank(publicKey)) {
            Map<String, String> map = generateKey();
            publicKey = map.get("PUB");
            privateKey = map.get("PRV");
            log.warn("SM2Processor is not configured with a key pair, use randomly generated key pair.\nPublicKey: {}\nPrivateKey: {}", publicKey, privateKey);
        }
        this.publicKey = publicKey;
        this.privateKey = privateKey;
        this.sm2 = SmUtil.sm2(privateKey, publicKey);
    }

    @Override
    public byte[] encrypt(byte[] data) {
       return sm2.encrypt(data, KeyType.PublicKey);
    }

    @Override
    public byte[] decrypt(String text) {
        return sm2.decrypt(text, KeyType.PrivateKey);
    }

    /**
     * 以base64编码生成秘钥对
     *
     * @return key="PUB"为公钥  key="PRV"为私钥
     */
    public static Map<String, String> generateKey() {
        return generateKey(SecurityMode.BASE64);
    }

    /**
     * 生成SM2秘钥对
     * map.get("PUB")：公钥
     * map.get("PRV")：私钥
     *
     * @param type {@link SecurityMode} 编码模式：type="hex"，其他默认base64
     * @return key="PUB"为公钥  key="PRV"为私钥
     */
    public static Map<String, String> generateKey(SecurityMode type) {
        Map<String, String> map = new HashMap<>();
        KeyPair pair = SecureUtil.generateKeyPair("SM2");
        String PRV;
        String PUB;
        switch (type) {
            case HEX:
                PRV = HexUtil.encodeHexStr(pair.getPrivate().getEncoded(), false);
                PUB = HexUtil.encodeHexStr(pair.getPublic().getEncoded(), false);
                break;
            case BASE64:
            default:
                PRV = Base64Encoder.encode(pair.getPrivate().getEncoded());
                PUB = Base64Encoder.encode(pair.getPublic().getEncoded());
        }
        map.put("PUB", PUB);
        map.put("PRV", PRV);
        return map;
    }

    public static void main(String[] args) {
        Map<String, String> map = generateKey(SecurityMode.HEX);
        //私钥
        String privateKey = map.get("PRV");
        System.out.println("privateKey = " + privateKey);
        //公钥
        String publicKey = map.get("PUB");
        System.out.println("publicKey = " + publicKey);

        SM2Processor sm2Processor = new SM2Processor(publicKey, privateKey);

        String encode = Base64Encoder.encode(sm2Processor.encrypt("Java".getBytes()));
        System.out.println("加密后：" + encode);
        byte[] decrypt = sm2Processor.decrypt(encode);
        System.out.println("解密后：" + StrUtil.str(decrypt, CharsetUtil.CHARSET_UTF_8));
    }
}
